import fileinput
import os

from neo4j import GraphDatabase
from config import NEO4J_CONFIG

# uri = "neo4j://192.168.88.161:7687" # 4.x版本用这个
uri = "bolt://192.168.88.161:7687" # 3.x版本用这个
driver = GraphDatabase.driver(uri, auth=("neo4j", "123456"), max_connection_lifetime=1000)
# driver = GraphDatabase.driver(**NEO4J_CONFIG)

def load_data(path:str):
    """
    加载审核之后疾病和症状对应的csv文件，生成一个字典
    :param path: 审核之后的csv所在目录
    :return: 字典，{"疾病名称":["症状名称1"...]}
    """
    disease_csv_list = os.listdir(path)
    disease_list = [csv.split('.')[0] for csv in disease_csv_list]

    symptom_list = []
    for disease_csv in disease_csv_list:
        print(disease_csv)
        symptom = list(map(lambda x: x.strip(),
                        fileinput.FileInput(os.path.join(path, disease_csv))))
        # print(disease_csv, symptom)
        symptom = list(filter(lambda x: 0<len(x)<100 ,symptom))

        symptom_list.append(symptom)

    return dict(zip(disease_list, symptom_list))


def write_to_neo4j(disease_symptom_dict:dict):
    """
    把传入的字典写入到neo4j中并创建关系
    :param disease_symptom_dict: 字典，{"疾病名称":["症状名称1"...]}
    :return:
    """
    with driver.session() as sess:
        for key, value in disease_symptom_dict.items():
            cypher = "MERGE (a:Disease{name:%r}) RETURN a" % key
            sess.run(cypher)
            for sym in value:
                cypher = "MERGE (b:Symptom{name:%r}) RETURN b" % sym
                sess.run(cypher)
                cypher = "MATCH (a:Disease{name:%r}) MATCH (b:Symptom{name:%r}) \
                                          WITH a,b MERGE(a)-[r:dis_to_sym]-(b)" % (key, sym)
                sess.run(cypher)

        cypher = "CREATE INDEX ON:Disease(name)"
        sess.run(cypher)
        cypher = "CREATE INDEX ON:Symptom(name)"
        sess.run(cypher)


if __name__ == '__main__':
    d = load_data("/root/ai17bj/doctor_offline/structured/reviewed")
    write_to_neo4j(d)
